import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import accuracy_score, classification_report
from sklearn.preprocessing import StandardScaler
from collections import Counter
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

print("加载数据...")
train_processed=pd.read_csv('1/output/MBA_processed.csv')
train_encoded=pd.read_csv('1/output/MBA_encoded.csv')
test_raw=pd.read_csv('1/test.csv')
# 预处理测试数据
def preprocess_test_data(df):
    df_processed=df.copy()
    if df_processed['race'].isnull().sum()>0:
        df_processed['race']=df_processed['race'].fillna(train_processed['race'].mode()[0])
    df_encoded=df_processed.copy()
    df_encoded['gender']=df_encoded['gender'].map({'Male':0,'Female':1})
    df_encoded['international']=df_encoded['international'].astype(int)
    for col in ['major','race','work_industry']:
        dummies=pd.get_dummies(df_encoded[col],prefix=col)
        df_encoded=df_encoded.drop(col,axis=1)
        df_encoded=pd.concat([df_encoded,dummies],axis=1)
    df_encoded=df_encoded.reindex(columns=train_encoded.columns,fill_value=0)
    return df_processed,df_encoded
test_processed,test_encoded=preprocess_test_data(test_raw)
# 准备数据
feature_cols=[col for col in train_encoded.columns if col not in ['admission']]
admission_mapping={'Reject':0,'Waitlist':1,'Admit':2}
X_train=train_encoded[feature_cols]
y_train=train_processed['admission'].map(admission_mapping)
X_test=test_encoded[feature_cols]
y_test=test_processed['admission'].map(admission_mapping)

# 2. 线性回归和逻辑回归分类
print("\n线性回归和逻辑回归分类")
# 线性回归分类（添加数据标准化）
scaler_linear=StandardScaler()
X_train_scaled_linear=scaler_linear.fit_transform(X_train)
X_test_scaled_linear=scaler_linear.transform(X_test)
linear_model=LinearRegression()
linear_model.fit(X_train_scaled_linear,y_train)
y_train_pred_linear_raw=linear_model.predict(X_train_scaled_linear)
y_test_pred_linear_raw=linear_model.predict(X_test_scaled_linear)
# 转换为分类结果
def linear_to_class(pred):
    result=np.zeros_like(pred,dtype=int)
    result[pred<0.5]=0
    result[(pred>=0.5)&(pred<1.5)]=1
    result[pred>=1.5]=2
    return result
y_train_pred_linear=linear_to_class(y_train_pred_linear_raw)
y_test_pred_linear=linear_to_class(y_test_pred_linear_raw)
train_acc_linear=accuracy_score(y_train,y_train_pred_linear)
test_acc_linear=accuracy_score(y_test,y_test_pred_linear)
print(f"线性回归 - 训练集: {train_acc_linear:.4f},测试集: {test_acc_linear:.4f}")
print("线性回归测试集分类报告:")
print(classification_report(y_test,y_test_pred_linear,target_names=['Reject','Waitlist','Admit'],zero_division=0))

# 逻辑回归分类（添加数据标准化）
scaler_logistic=StandardScaler()
X_train_scaled=scaler_logistic.fit_transform(X_train)
X_test_scaled=scaler_logistic.transform(X_test)
logistic_model=LogisticRegression(random_state=42,max_iter=2000,class_weight='balanced',solver='lbfgs')
logistic_model.fit(X_train_scaled,y_train)
y_train_pred_logistic=logistic_model.predict(X_train_scaled)
y_test_pred_logistic=logistic_model.predict(X_test_scaled)
train_acc_logistic=accuracy_score(y_train,y_train_pred_logistic)
test_acc_logistic=accuracy_score(y_test,y_test_pred_logistic)
print(f"逻辑回归 - 训练集: {train_acc_logistic:.4f},测试集: {test_acc_logistic:.4f}")
print("逻辑回归测试集分类报告:")
print(classification_report(y_test,y_test_pred_logistic,target_names=['Reject','Waitlist','Admit'],zero_division=0))

# 3. sklearn MLP分类
print("\nsklearn MLP分类")
scaler_mlp=StandardScaler()
X_train_scaled_mlp=scaler_mlp.fit_transform(X_train)
X_test_scaled_mlp=scaler_mlp.transform(X_test)
mlp_sklearn=MLPClassifier(hidden_layer_sizes=(128,256),activation='relu',
                          learning_rate_init=0.001,max_iter=500,batch_size=64,
                          random_state=42,early_stopping=True,solver='adam')
mlp_sklearn.fit(X_train_scaled_mlp,y_train)
y_train_pred_mlp=mlp_sklearn.predict(X_train_scaled_mlp)
y_test_pred_mlp=mlp_sklearn.predict(X_test_scaled_mlp)
train_acc_mlp=accuracy_score(y_train,y_train_pred_mlp)
test_acc_mlp=accuracy_score(y_test,y_test_pred_mlp)
print(f"sklearn MLP - 训练集: {train_acc_mlp:.4f},测试集: {test_acc_mlp:.4f}")
print("sklearn MLP测试集分类报告:")
print(classification_report(y_test,y_test_pred_mlp,target_names=['Reject','Waitlist','Admit'],zero_division=0))

# 4. PyTorch MLP分类
print("\nPyTorch MLP分类")
# 定义MLP模型
def create_mlp(input_dim,output_dim=3):
    model=nn.Sequential(
        nn.Linear(input_dim,128),
        nn.ReLU(),
        nn.Linear(128,256),
        nn.ReLU(),
        nn.Linear(256,output_dim)
    )
    return model
# 数据准备（使用标准化数据）
X_train_tensor=torch.FloatTensor(X_train_scaled_mlp.astype(float))
X_test_tensor=torch.FloatTensor(X_test_scaled_mlp.astype(float))
y_train_tensor=torch.LongTensor(y_train.astype(int).values)
y_test_tensor=torch.LongTensor(y_test.astype(int).values)
# 创建数据加载器
train_dataset=TensorDataset(X_train_tensor,y_train_tensor)
train_loader=DataLoader(train_dataset,batch_size=64,shuffle=True)
# 模型初始化
input_dim=X_train.shape[1]
model=create_mlp(input_dim)
# 计算权重
total_samples=len(y_train)
reject_count=sum(y_train==0)
waitlist_count=sum(y_train==1)
admit_count=sum(y_train==2)
weight_reject=min(total_samples/(3*reject_count),2.0)
weight_waitlist=min(total_samples/(3*waitlist_count),50.0)
weight_admit=min(total_samples/(3*admit_count),10.0)
class_weights=torch.FloatTensor([weight_reject,weight_waitlist,weight_admit])
criterion=nn.CrossEntropyLoss(weight=class_weights)
optimizer=optim.Adam(model.parameters(),lr=0.001)
# 训练
model.train()
for epoch in range(100):
    total_loss=0
    for batch_X,batch_y in train_loader:
        optimizer.zero_grad()
        outputs=model(batch_X)
        loss=criterion(outputs,batch_y)
        loss.backward()
        optimizer.step()
        total_loss+=loss.item()
    if (epoch+1)%20==0:
        avg_loss=total_loss/len(train_loader)
        print(f"Epoch [{epoch+1}/100],平均损失: {avg_loss:.4f}")
# 模型评估
model.eval()
with torch.no_grad():
    train_outputs=model(X_train_tensor)
    _,train_predicted=torch.max(train_outputs,1)
    train_acc_pytorch=(train_predicted==y_train_tensor).float().mean()
    test_outputs=model(X_test_tensor)
    _,test_predicted=torch.max(test_outputs,1)
    test_acc_pytorch=(test_predicted==y_test_tensor).float().mean()
print(f"PyTorch MLP - 训练集: {train_acc_pytorch:.4f},测试集: {test_acc_pytorch:.4f}")
print("PyTorch MLP测试集分类报告:")
test_predicted_numpy=test_predicted.cpu().numpy()
print(classification_report(y_test,test_predicted_numpy,target_names=['Reject','Waitlist','Admit'],zero_division=0))

# 5. 激活函数对比
print("\n激活函数对比")
activations=['identity','relu','logistic']
activation_results={}
for act in activations:
    mlp_act=MLPClassifier(hidden_layer_sizes=(128,256),activation=act,
                         learning_rate_init=0.001,max_iter=500,batch_size=64,
                         random_state=42,early_stopping=True,solver='adam')
    mlp_act.fit(X_train_scaled_mlp,y_train)
    y_test_pred_act=mlp_act.predict(X_test_scaled_mlp)
    test_acc_act=accuracy_score(y_test,y_test_pred_act)
    activation_results[act]=test_acc_act
    act_name='sigmoid' if act=='logistic' else act
    print(f"{act_name}激活函数 - 测试集: {test_acc_act:.4f}")
    print(f"{act_name}激活函数分类报告:")
    print(classification_report(y_test,y_test_pred_act,target_names=['Reject','Waitlist','Admit'],zero_division=0))

# 6. 特征工程
print("\n特征工程")
# 创建新特征
X_engineered=pd.DataFrame()
X_engineered['gpa']=train_processed['gpa']
X_engineered['gmat']=train_processed['gmat']
X_engineered['work_exp']=train_processed['work_exp']
X_engineered['gpa_gmat_interaction']=train_processed['gpa']*train_processed['gmat']
X_engineered['gmat_high']=(train_processed['gmat']>=700).astype(int)
X_engineered['gmat_low']=(train_processed['gmat']<650).astype(int)
X_engineered['is_business']=(train_processed['major']=='Business').astype(int)
X_engineered['is_stem']=(train_processed['major']=='STEM').astype(int)
X_engineered['is_international']=train_processed['international'].astype(int)
X_engineered['is_female']=(train_processed['gender']=='Female').astype(int)
# 标准化
scaler=StandardScaler()
X_engineered_scaled=scaler.fit_transform(X_engineered)
# 对测试集应用相同特征工程
X_test_engineered=pd.DataFrame()
X_test_engineered['gpa']=test_processed['gpa']
X_test_engineered['gmat']=test_processed['gmat']
X_test_engineered['work_exp']=test_processed['work_exp']
X_test_engineered['gpa_gmat_interaction']=test_processed['gpa']*test_processed['gmat']
X_test_engineered['gmat_high']=(test_processed['gmat']>=700).astype(int)
X_test_engineered['gmat_low']=(test_processed['gmat']<650).astype(int)
X_test_engineered['is_business']=(test_processed['major']=='Business').astype(int)
X_test_engineered['is_stem']=(test_processed['major']=='STEM').astype(int)
X_test_engineered['is_international']=test_processed['international'].astype(int)
X_test_engineered['is_female']=(test_processed['gender']=='Female').astype(int)
X_test_engineered_scaled=scaler.transform(X_test_engineered)
# 特征工程MLP
mlp_optimized=MLPClassifier(hidden_layer_sizes=(128,256),activation='relu',
                            learning_rate_init=0.001,max_iter=500,batch_size=64,
                            random_state=42,early_stopping=True,solver='adam')
mlp_optimized.fit(X_engineered_scaled,y_train)
y_eng_train_pred=mlp_optimized.predict(X_engineered_scaled)
y_eng_test_pred=mlp_optimized.predict(X_test_engineered_scaled)
train_acc_eng=accuracy_score(y_train,y_eng_train_pred)
test_acc_eng=accuracy_score(y_test,y_eng_test_pred)
print(f"特征工程MLP - 训练集: {train_acc_eng:.4f},测试集: {test_acc_eng:.4f}")
print("特征工程MLP测试集分类报告:")
print(classification_report(y_test,y_eng_test_pred,target_names=['Reject','Waitlist','Admit'],zero_division=0))

# 7. 结果总结
print("\n" + "=" * 60)
print("分类任务结果总结")
print("=" * 60)

results_summary={
    '线性回归':{'训练集':train_acc_linear,'测试集':test_acc_linear},
    '逻辑回归':{'训练集':train_acc_logistic,'测试集':test_acc_logistic},
    'sklearn MLP':{'训练集':train_acc_mlp,'测试集':test_acc_mlp},
    'PyTorch MLP':{'训练集':train_acc_pytorch.item(),'测试集':test_acc_pytorch.item()},
    'MLP(特征工程)':{'训练集':train_acc_eng,'测试集':test_acc_eng}
}

print("\n模型性能对比:")
print("-" * 50)
print(f"{'模型名称':<15} | {'训练集':<8} | {'测试集':<8}")
print("-" * 50)
for model_name,scores in results_summary.items():
    print(f"{model_name:<15} | {scores['训练集']:<8.4f} | {scores['测试集']:<8.4f}")

print(f"\n激活函数对比结果:")
for act,acc in activation_results.items():
    act_name='sigmoid' if act=='logistic' else act
    print(f"{act_name:10}: {acc:.4f}")

print("\n分析完成!")